#include <stdio.h>
#include <pthread.h>
#include <string.h>
#include <stdlib.h>

#include "zmalloc.h"

#define PREFIX_SIZE (sizeof(size_t))

static size_t used_memory = 0;
static int zmalloc_thread_safe = 0;
pthread_mutex_t used_memory_mutex = PTHREAD_MUTEX_INITIALIZER;

#define update_zmalloc_state_alloc(__n, __size) do { \
	size_t _n = (__n); \
	if(_n&(sizeof(long)-1)) _n += sizeof(long) - (_n&(sizeof(long)-1)); \
	printf("alloc : __n:%d,_n:%d\n", __n, _n);\
	if(zmalloc_thread_safe) { \
		pthread_mutex_lock(&used_memory_mutex); \
		used_memory += _n; \
		pthread_mutex_unlock(&used_memory_mutex); \
	} else { \
		used_memory += _n; \
	} \
}while(0)


#define update_zmalloc_state_free(__n) do { \
	size_t _n = (__n); \
	if(_n&(sizeof(long)-1)) _n += sizeof(long) - (_n&(sizeof(long)-1)); \
	printf("free : __n:%d, _n:%d\n", __n, _n);\
	if(zmalloc_thread_safe) { \
		pthread_mutex_lock(&used_memory_mutex); \
		used_memory -= _n; \
		pthread_mutex_unlock(&used_memory_mutex); \
	} else { \
		used_memory -= _n; \
	} \
}while(0)


static void zmalloc_oops(size_t size)
{
	fprintf(stderr, "[%s]:zmalloc: Out of memory trying to allocate %zu \
			bytes\n", __FUNCTION__, size);
	fflush(stderr);
	abort();
}

void *zmalloc(size_t size)
{
	void *p = malloc(size+PREFIX_SIZE);
	if(p==NULL)
		zmalloc_oops(size);

	*((size_t*)p) = size;
	update_zmalloc_state_alloc(size+PREFIX_SIZE, size);
	return (char*)p+PREFIX_SIZE;
}

void *zcalloc(size_t size)
{
	void *p = calloc(1, size+PREFIX_SIZE);
	if(p==NULL)
		zmalloc_oops(size);
	update_zmalloc_state_alloc(size+PREFIX_SIZE, size);
	return (char*)p + PREFIX_SIZE;
}

void *zrealloc(void *ptr, size_t size)
{
	if(ptr==NULL)
		return zmalloc(size);

	void *realptr = (char*)ptr-PREFIX_SIZE;
	size_t oldsize = *((size_t*)realptr);
	update_zmalloc_state_free(oldsize);

	void *newptr = zmalloc(size);
	if(newptr==NULL)
		zmalloc_oops(size);

	update_zmalloc_state_alloc(size+PREFIX_SIZE, size);
	return (char*)newptr+PREFIX_SIZE;
}

void zfree(void *ptr)
{
	if(ptr==NULL)
		return;

	void *realptr = (char*)ptr-PREFIX_SIZE;
	size_t free_size = *((size_t*)realptr);
	printf("[%d]:%s free size:%d\n", getpid(), __FUNCTION__, free_size);
	update_zmalloc_state_free(free_size+PREFIX_SIZE);
	free(realptr);
	return;
}

char *zstrdup(const char *s)
{
	size_t len = strlen(s)+1;
	void *p = zmalloc(len);
	memcpy(p, s, len);
	return p;
}

size_t zmalloc_used_memory(void)
{
	size_t um;
	if(zmalloc_thread_safe)
		pthread_mutex_lock(&used_memory_mutex);
	um = used_memory;
	if(zmalloc_thread_safe)
		pthread_mutex_unlock(&used_memory_mutex);
	return um;
}

void zmalloc_enable_thread_safeness(void)
{
	zmalloc_thread_safe = 1;
}

size_t zmalloc_get_rss(void)
{
	return zmalloc_used_memory();
}
